"""
1.	利用Pytorch深度学习框架，根据yolov2的主干网络darknet19模型，进行猫狗数据集的模型训练和模型推理，按下面要求完成相应代码。（76分）
"""
import torch as pt
import numpy as np
import matplotlib.pyplot as plt
import os
import sys
import cv2 as cv
from sklearn.model_selection import train_test_split
from torch.utils.data import TensorDataset, DataLoader

DATA_DIR = 'data_zoo'
PIC_SIZE = 224
LEARNING_RATE = 0.001
BATCH_SIZE = 16
N_EPOCHS = 8

# Use GPU if it is available
device = 'cuda:0' if pt.cuda.is_available() else 'cpu'
print('device:', device)
device = pt.device(device)

# ①	读入猫狗数据集（data_zoo），并进行必要的预处理
x, y = [], []
lable_value = 0
print('Loading...')
for sub_dir in os.listdir(DATA_DIR):
    sub_dir_path = os.path.join(DATA_DIR, sub_dir)
    print(f'Loading {sub_dir}')
    for file_name in os.listdir(sub_dir_path):
        file_path = os.path.join(sub_dir_path, file_name)
        img = cv.imread(file_path, cv.IMREAD_COLOR)
        img = cv.resize(img, (PIC_SIZE, PIC_SIZE), interpolation=cv.INTER_CUBIC)
        img = img.astype(np.float32) / 255. * 2. - 1.
        x.append(img)
        y.append(lable_value)
    print(f'Loaded {sub_dir}')
    lable_value += 1
print('Loaded')
x = np.float32(x)
y = np.int64(y)

# ②	划分数据为训练集和测试集（比例自拟）
x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.8, random_state=1, shuffle=True)

# ③	在训练集中随机选取9张图片进行3行3列图像显示
print('Showing some data, check and close the plotting window to go on ...')
spr = 3
spc = 3
spn = 0
plt.figure(figsize=[6, 6])
for i in range(9):
    img = x_train[i]
    img += 1.
    img /= 2.
    img = img[:, :, ::-1]
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(str(y_train[i]))
    plt.imshow(img)
    plt.axis('off')
plt.show()


# ④	构建卷积单元类封装ConvCell，由卷积、BN、Relu激活组成
class ConvCell(pt.nn.Module):

    def __init__(self, in_ch, out_ch, ksize=(3, 3), stride=(1, 1), padding=1, **kwargs):
        super().__init__(**kwargs)
        self.conv = pt.nn.Conv2d(in_ch, out_ch, ksize, stride, padding)
        self.bn = pt.nn.BatchNorm2d(out_ch)
        self.relu = pt.nn.ReLU()

    def forward(self, inputs):
        x = self.conv(inputs)
        x = self.bn(x)
        x = self.relu(x)
        return x


# ⑤	参见附图，构建yolov2的主干网络类封装DarkNet19
class DarkNet19(pt.nn.Module):

    def __init__(self, n_cls, **kwargs):
        super().__init__(**kwargs)
        self.seq = pt.nn.Sequential(
            ConvCell(3, 32),
            pt.nn.MaxPool2d((2, 2), (2, 2)),

            ConvCell(32, 64),
            pt.nn.MaxPool2d((2, 2), (2, 2)),

            ConvCell(64, 128),
            ConvCell(128, 64, (1, 1), (1, 1), 0),
            ConvCell(64, 128),
            pt.nn.MaxPool2d((2, 2), (2, 2)),

            ConvCell(128, 256),
            ConvCell(256, 128, (1, 1), (1, 1), 0),
            ConvCell(128, 256),
            pt.nn.MaxPool2d((2, 2), (2, 2)),

            ConvCell(256, 512),
            ConvCell(512, 256, (1, 1), (1, 1), 0),
            ConvCell(256, 512),
            ConvCell(512, 256, (1, 1), (1, 1), 0),
            ConvCell(256, 512),
            pt.nn.MaxPool2d((2, 2), (2, 2)),

            ConvCell(512, 1024),
            ConvCell(1024, 512, (1, 1), (1, 1), 0),
            ConvCell(512, 1024),
            ConvCell(1024, 512, (1, 1), (1, 1), 0),
            ConvCell(512, 1024),

            ConvCell(1024, n_cls, (1, 1), (1, 1), 0),
            pt.nn.AdaptiveAvgPool2d(1),
        )

    # ⑥	进行前向传播
    def forward(self, inputs):
        x = self.seq(inputs)
        x = pt.squeeze(x, 3)
        x = pt.squeeze(x, 2)
        return x


# ⑦	打印输出整体模型维度结构
model = DarkNet19(2).to(device)
print(model)

# ⑧	正确合理使用损失函数和优化器
loss = pt.nn.CrossEntropyLoss()
optim = pt.optim.Adam(model.parameters(), lr=LEARNING_RATE)


def accuracy(y, pred):
    pred = pred.argmax(dim=1)
    acc = pt.eq(y, pred).float().mean()
    return acc


# ⑨	训练集数据进行模型训练，自拟合适的超参数
x_train = np.transpose(x_train, (0, 3, 1, 2))
x_train = pt.Tensor(x_train)
y_train = pt.Tensor(y_train)
ds_train = TensorDataset(x_train, y_train)
dl_train = DataLoader(ds_train, BATCH_SIZE)

x_test = np.transpose(x_test, (0, 3, 1, 2))
x_test = pt.Tensor(x_test)
y_test = pt.Tensor(y_test)
ds_test = TensorDataset(x_test, y_test)
dl_test = DataLoader(ds_test, BATCH_SIZE)


def process_data(dl, label, is_train):
    n_batch = len(dl)
    group = int(np.ceil(n_batch / 10))
    avg_loss, avg_acc = 0., 0.
    for i, (bx, by) in enumerate(dl):
        bx = bx.to(device)
        by = by.long().to(device)
        if is_train:
            model.train(True)
            optim.zero_grad()
            pred = model(bx)
            lossv = loss(pred, by)
            lossv.backward()
            optim.step()
            model.train(False)
        else:
            model.eval()
            pred = model(bx)
            lossv = loss(pred, by)
        accv = accuracy(by, pred)
        lossv = lossv.detach().cpu().numpy()
        accv = accv.detach().cpu().numpy()
        avg_loss += lossv
        avg_acc += accv
        if i % group == 0 or i == n_batch - 1:
            print(f'{label}: #{i + 1}: loss = {lossv}, acc = {accv}')
    avg_loss /= i + 1
    avg_acc /= i + 1
    return avg_loss, avg_acc


# ⑩	模型训练完毕后，打印输出训练集准确率
loss_his, acc_his, loss_his_val, acc_his_val = [], [], [], []
print('Training')
for epoch in range(N_EPOCHS):
    avg_loss, avg_acc = process_data(dl_train, 'train', True)
    loss_his.append(avg_loss)
    acc_his.append(avg_acc)
    avg_loss_val, avg_acc_val = process_data(dl_test, 'val', False)
    loss_his_val.append(avg_loss_val)
    acc_his_val.append(avg_acc_val)
    print(f'Epoch#{i + 1}: loss = {avg_loss}, acc = {avg_acc}, loss_val = {avg_loss_val}, acc_val = {avg_acc_val}')
print('Trained')

# 11	打印输出测试集准确率
print('Testing ...')
avg_loss, avg_acc = process_data(dl_test, 'test', False)
print(f'Test: loss = {avg_loss}, acc = {avg_acc}')
print('Tested')

# plot
print('Plotting, check and close plotting window to terminate process ...')
spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])

spn += 1
plt.subplot(spr, spc, spn)
plt.title('Loss')
plt.plot(loss_his, label='train')
plt.plot(loss_his_val, label='val')
plt.legend()

spn += 1
plt.subplot(spr, spc, spn)
plt.title('Accuracy')
plt.plot(acc_his, label='train')
plt.plot(acc_his_val, label='val')
plt.legend()

plt.show()
print('Over')
